import heapq


def heap_insert(arr, idx):
    while arr[idx] > arr[int((idx-1) / 2)]:
        arr[idx], arr[int((idx-1) / 2)] = arr[int((idx-1) / 2)], arr[idx]
        idx = int((idx-1) / 2)


def heapify(arr, idx, size):
    left = 2 * idx + 1
    while left < size:
        largest = left + 1 if (left+1 < size and arr[left] < arr[left+1]) else left
        largest = largest if arr[largest] > arr[idx] else idx
        if largest == idx:
            break
        arr[idx], arr[largest] = arr[largest], arr[idx]
        idx = largest
        left = 2 * idx + 1


def heap_sort(arr):
    if not arr or len(arr) < 2:
        return
    for i in range(len(arr)):
        heap_insert(arr, i)
    heap_size = len(arr) - 1
    arr[0], arr[heap_size] = arr[heap_size], arr[0]
    while heap_size > 0:
        heapify(arr, 0, heap_size)
        heap_size -= 1
        arr[0], arr[heap_size] = arr[heap_size], arr[0]


# 例子
def sort_arr_less_k(arr, K):
    """一个几乎有序的数组，每个元素排序后移动的距离不超过K"""
    heap = []
    for idx in range(min(K, len(arr))):
        heapq.heappush(heap, arr[idx])
    right = min(K, len(arr))
    idx = 0
    while right < len(arr):
        arr[idx] = heapq.heappop(heap)
        heapq.heappush(heap, arr[right])
        right += 1
        idx += 1
    while heap:
        arr[idx] = heapq.heappop(heap)
        idx += 1


if __name__ == "__main__":
    arr = [4, 1, 6, 2, 9, 5]
    sort_arr_less_k(arr, 3)
    print(arr)
